package cn.itcast.flink.trans;

import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * Author itcast
 * Date 2022/1/13 10:18
 * Desc 生成一个随机数字 1~100，过滤出来 90个 数字
 * 先不重分布，统计一下每个taskid(分区id,cpuid)，处理数字(1)
 * 对taskid进行分组，统计
 * 打印输出
 */
public class RebalanceDemo {
    public static void main(String[] args) throws Exception {
        //todo 1.获取流执行环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //todo 2.设置并行度为 3
        env.setParallelism(3);
        //todo 3.生成1~100的序列，过滤出来大于10的数字
        SingleOutputStreamOperator<Long> source = env.fromSequence(1, 100)
                .filter(new FilterFunction<Long>() {
            @Override
            public boolean filter(Long value) throws Exception {
                return value > 10;
            }
        });
        //todo 4.将序列转换map Tuple2[Taskid,1] 通过上下文context
       /* SingleOutputStreamOperator<Tuple2<Integer, Integer>> mapStream = source
                .map(new RichMapFunction<Long, Tuple2<Integer*//*taskid*//*, Integer*//*1*//*>>() {
            @Override
            public Tuple2<Integer, Integer> map(Long value) throws Exception {
                //通过当前运行时上下文获取数据被处理的 cpu index 的id
                int taskid = getRuntimeContext().getIndexOfThisSubtask();
                return Tuple2.of(taskid, 1);
            }
        });*/
        //todo 将数据进行重分布，数据在下游subtask中均匀分布
        SingleOutputStreamOperator<Tuple2<Integer, Integer>> mapStream = source
                .rebalance()
                .map(new RichMapFunction<Long, Tuple2<Integer/*taskid*/, Integer/*1*/>>() {
                    @Override
                    public Tuple2<Integer, Integer> map(Long value) throws Exception {
                        //通过当前运行时上下文获取数据被处理的 cpu index 的id
                        int taskid = getRuntimeContext().getIndexOfThisSubtask();
                        return Tuple2.of(taskid, 1);
                    }
                });

        //todo 5.根据 Taskid 分组
        SingleOutputStreamOperator<Tuple2<Integer, Integer>> result = mapStream.keyBy(t -> t.f0)
                //todo 6.聚合
                .sum(1);
        //todo 7.打印输出
        result.print();
        //todo 8.执行流环境
        env.execute();
    }
}
